// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// A command line executable that generates a bunch of valid IPC files
// containing example tensors.  Those are used as fuzzing seeds to make
// fuzzing more efficient.

#include <cstdlib>
#include <iostream>
#include <memory>
#include <string>
#include <vector>

#include <nebula/io/file.h>
#include <nebula/io/memory.h>
#include <tests/ipc/test_common.h>
#include <nebula/ipc/writer.h>

#include <nebula/types/tensor.h>
#include <turbo/files//io_util.h>

namespace nebula {
namespace ipc {

turbo::Result<turbo::FilePath> PrepareDirectory(const std::string& dir) {
  TURBO_MOVE_OR_RAISE(auto dir_fn, turbo::FilePath::from_string(dir));
  TURBO_RETURN_NOT_OK(::turbo::create_directory(dir_fn));
  return std::move(dir_fn);
}

turbo::Result<std::shared_ptr<Buffer>> MakeSerializedBuffer(
    std::function<turbo::Status(const std::shared_ptr<io::BufferOutputStream>&)> fn) {
  TURBO_MOVE_OR_RAISE(auto sink, io::BufferOutputStream::create(1024));
  TURBO_RETURN_NOT_OK(fn(sink));
  return sink->finish();
}

turbo::Result<std::shared_ptr<Buffer>> SerializeTensor(const std::shared_ptr<Tensor>& tensor) {
  return MakeSerializedBuffer(
      [&](const std::shared_ptr<io::BufferOutputStream>& sink) -> turbo::Status {
        int32_t metadata_length;
        int64_t body_length;
        return ipc::write_tensor(*tensor, sink.get(), &metadata_length, &body_length);
      });
}

turbo::Result<std::vector<std::shared_ptr<Tensor>>> Tensors() {
  std::vector<std::shared_ptr<Tensor>> tensors;
  std::shared_ptr<Tensor> tensor;
  std::vector<int64_t> shape = {5, 3, 7};
  std::shared_ptr<DataType> types[] = {int8(),  int16(),  int32(),  int64(),
                                       uint8(), uint16(), uint32(), uint64()};
  uint32_t seed = 0;
  for (auto type : types) {
    TURBO_RETURN_NOT_OK(
        test::MakeRandomTensor(type, shape, /*row_major_p=*/true, &tensor, seed++));
    tensors.push_back(tensor);
    TURBO_RETURN_NOT_OK(
        test::MakeRandomTensor(type, shape, /*row_major_p=*/false, &tensor, seed++));
    tensors.push_back(tensor);
  }
  return tensors;
}

turbo::Status GenerateTensors(const turbo::FilePath& dir_fn) {
  int sample_num = 1;
  auto sample_name = [&]() -> std::string {
    return "tensor-" + std::to_string(sample_num++);
  };

  TURBO_MOVE_OR_RAISE(auto tensors, Tensors());

  for (const auto& tensor : tensors) {
    TURBO_MOVE_OR_RAISE(auto buf, SerializeTensor(tensor));
    auto sample_fn =  dir_fn/(sample_name());
    std::cerr << sample_fn.string() << std::endl;
    TURBO_MOVE_OR_RAISE(auto file, io::FileOutputStream::open(sample_fn.string()));
    TURBO_RETURN_NOT_OK(file->write(buf));
    TURBO_RETURN_NOT_OK(file->close());
  }
  return turbo::OkStatus();
}

turbo::Status DoMain(const std::string& out_dir) {
  TURBO_MOVE_OR_RAISE(auto dir_fn, PrepareDirectory(out_dir));
  return GenerateTensors(dir_fn);
}

TURBO_NORETURN void Usage() {
  std::cerr << "Usage: nebula-ipc-generate-tensor-fuzz-corpus "
            << "-stream <output directory>" << std::endl;
  std::exit(2);
}

int Main(int argc, char** argv) {
  if (argc != 3) {
    Usage();
  }

  auto opt = std::string(argv[1]);
  if (opt != "-stream") {
    Usage();
  }

  auto out_dir = std::string(argv[2]);

  turbo::Status st = DoMain(out_dir);
  if (!st.ok()) {
    std::cerr << st.to_string() << std::endl;
    return 1;
  }
  return 0;
}

}  // namespace ipc
}  // namespace nebula

int main(int argc, char** argv) { return nebula::ipc::Main(argc, argv); }
